#include <atomic>
#include <chrono>
#include <cstdint>
#include <thread>
#include <stdexcept>
#include <iostream>
#include <vector>

class FixedWindowLimiter {
public:
	using clock = std::chrono::steady_clock;

	// window_size_ms: 窗口大小（毫秒）
	// capacity: 每窗口允许的最大请求数
	FixedWindowLimiter(uint64_t window_size_ms, uint64_t capacity)
		: window_size_ms_(window_size_ms),
		  capacity_(capacity),
		  window_id_(0),
		  count_(0)
	{
		if (window_size_ms_ == 0) throw std::invalid_argument("window_size_ms must be > 0");
		if (capacity_ == 0) throw std::invalid_argument("capacity must be > 0");
		// 初始化 window_id_ 为当前窗口，避免启动时的异常涨落
		window_id_.store(current_window(), std::memory_order_relaxed);
	}

	// 非阻塞尝试获取 n 个配额；成功返回 true，失败返回 false（立即返回）
	bool try_acquire(uint64_t n = 1)
	{
		if (n == 0) return true;
		auto cur = current_window();

		for (;;) {
			uint64_t w = window_id_.load(std::memory_order_acquire);
			if (w == cur) {
				// 同一窗口，原子加
				uint64_t prev = count_.fetch_add(n, std::memory_order_acq_rel);
				if (prev + n <= capacity_) {
					return true;
				}
				else {
					// 超额，回退并失败
					count_.fetch_sub(n, std::memory_order_acq_rel);
					return false;
				}
			}
			else if (w < cur) {
				// 窗口已过期，尝试推进到当前窗口（CAS）
				// 只有推进窗口成功的线程负责初始化 count_ 为 n
				if (window_id_.compare_exchange_strong(w, cur,
				                                       std::memory_order_acq_rel,
				                                       std::memory_order_acquire)) {
					// 我们成功把 window_id 推进到 cur，初始化 count_
					count_.store(n, std::memory_order_release);
					// 检查是否超过 capacity
					return n <= capacity_;
				}
				// CAS 失败，说明其他线程推进了 window_id，重试
			}
			else {
				// w > cur (极少见，理论上不应发生，因为 steady_clock 递增)
				// 为了健壮性，简单重试（也可以当作拒绝）
				// 这里使用 yield 避免 busy-wait 太激烈
				std::this_thread::yield();
			}

			// 更新 cur（时间可能已跨窗口）
			cur = current_window();
		}
	}

	// 阻塞获取 n 个配额（会在无法立即拿到时短暂 sleep，直到可用）
	// 返回实际等待的毫秒数（可能为 0）
	uint64_t acquire(uint64_t n = 1)
	{
		if (n == 0) return 0;
		while (true) {
			if (try_acquire(n)) return 0;
			// 计算到下一个窗口的剩余时间并睡眠一段时间（避免忙等）
			uint64_t ms_to_next = ms_until_next_window();
			// sleep 最少 1ms，最多 ms_to_next（防止 0 导致 busy-loop）
			if (ms_to_next == 0) {
				std::this_thread::yield();
			}
			else {
				std::this_thread::sleep_for(std::chrono::milliseconds(std::min<uint64_t>(ms_to_next, 5)));
			}
		}
	}

	// 近似读取当前窗口已使用的配额（仅用于监控，非精确保证）
	uint64_t approx_count() const
	{
		return count_.load(std::memory_order_relaxed);
	}

	uint64_t capacity() const
	{
		return capacity_;
	}
	uint64_t window_size_ms() const
	{
		return window_size_ms_;
	}

private:
	uint64_t current_window() const
	{
		auto now = clock::now();
		auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count();
		return static_cast<uint64_t>(ms / window_size_ms_);
	}

	uint64_t ms_until_next_window() const
	{
		auto now = clock::now();
		auto ms = std::chrono::duration_cast<std::chrono::milliseconds>(now.time_since_epoch()).count();
		uint64_t curWinStartMs = (ms / window_size_ms_) * window_size_ms_;
		uint64_t nextStartMs = curWinStartMs + window_size_ms_;
		if (nextStartMs <= static_cast<uint64_t>(ms)) return 0;
		return nextStartMs - static_cast<uint64_t>(ms);
	}

private:
	const uint64_t window_size_ms_;
	const uint64_t capacity_;

	// 注意：这两个原子变量决不能在外部以非原子方式同时访问
	std::atomic<uint64_t> window_id_;
	std::atomic<uint64_t> count_;
};

// ======================= 使用示例 =======================
int main()
{
	FixedWindowLimiter limiter(1000 /* 1s 窗口 */, 10 /* 每窗口最多 10 次 */);

	constexpr int kThreads = 8;
	constexpr int kAttemptsPerThread = 20;

	std::vector<std::thread> ths;
	std::atomic<int> success{0}, fail{0};

	for (int t = 0; t < kThreads; ++t) {
		ths.emplace_back([&]() {
			for (int i = 0; i < kAttemptsPerThread; ++i) {
				if (limiter.try_acquire()) {
					success.fetch_add(1, std::memory_order_relaxed);
					// 模拟短暂工作
					std::this_thread::sleep_for(std::chrono::milliseconds(10));
				}
				else {
					fail.fetch_add(1, std::memory_order_relaxed);
					// 不阻塞，稍后再试
					std::this_thread::sleep_for(std::chrono::milliseconds(5));
				}
			}
		});
	}

	for (auto& th : ths) th.join();

	std::cout << "Success: " << success.load() << ", Fail: " << fail.load() << "\n";
	return 0;
}
